#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Time   : 2022/9/8 19:10
# @Author : cjw
from typing import Optional, Union, List

from pymysql import connect, Error, Connection
from pymysql.cursors import DictCursor, Cursor
from loguru import logger


class PyMySQL:

    def __init__(self, host: str, user: str, password: str, database: str,
                 port: int = 3306, is_dict_cursor: bool = True, **kwargs):
        """
        初始化
        :param host: 数据库服务地址
        :param user: 用户名
        :param password: 密码
        :param database: 需要连接的数据库
        :param port: 数据库端口
        :param is_dict_cursor: 是否将查询结果转换为字典
        :param kwargs:
        """
        self._host = host
        self._user = user
        self._password = password
        self._database = database
        self._port = port
        self._is_dict_cursor = is_dict_cursor
        self._kwargs = kwargs
        self.connection = self.get_connection()

    def get_connection(self):
        """获取数据库连接"""
        try:
            conn: Connection = connect(
                host=self._host,
                user=self._user,
                password=self._password,
                database=self._database,
                port=self._port,
                cursorclass=DictCursor if self._is_dict_cursor else Cursor,
                **self._kwargs
            )
            return conn
        except Error as e:
            logger.error(f'连接失败，异常信息：{e}')
            return None

    def close(self):
        """关闭数据库连接"""
        try:
            if self.connection:
                self.connection.close()
            else:
                logger.info('未初始化数据库连接，不能关闭')
        except Exception as e:
            logger.error(f'数据库关闭失败，报错信息：{e}')

    def select(self, sql: str, params: Optional[Union[tuple, list, dict]] = None,
               rows: Union[int, None] = None) -> list:
        """
        数据库查询操作
        :param sql: 查询sql
        :param params: 查询参数
        :param rows: 默认返回所有数据
        :return:
        """
        if not self.connection:
            logger.error(msg := '数据库连接未初始化！')
            raise ValueError(msg)

        with self.connection.cursor() as cur:
            # 如果参数有?, 全部替换成%s
            cur.execute(sql.replace('?', '%s'), params)
            rs = cur.fetchmany(rows) if rows else cur.fetchall()
        return rs

    def execute(self, sql: str, params: Optional[List[tuple]] = None) -> int:
        """
        数据库执行增、删、改操作
        :param sql: sql语句
        :param params: sql语句参数，执行多条语句，参数格式：[(),()]
        :return: 返回sql执行影响的数据表的行数
        """
        if not self.connection:
            logger.error(msg := '数据库连接未初始化！')
            raise ValueError(msg)

        if not isinstance(params, list):
            logger.error(msg := f'参数格式不正确：{params}, 正确格式：[(),()]')
            raise ValueError(msg)

        with self.connection.cursor() as cur:
            if len(params) > 1:
                # 首先对sql语句中的?进行处理
                cur.executemany(sql.replace('?', '%s'), params)
            else:
                cur.execute(sql.replace('?', '%s'), *params if params else None)
            rowcount = cur.rowcount
        self.connection.commit()
        return rowcount

    def __del__(self):
        """
        销毁对象时，关闭数据库连接。
        备注：日志打印不起作用，后续调整
        :return:
        """
        try:
            if not self.connection:
                self.connection.close()
            else:
                logger.info('数据库连接未初始化，不需要关闭!')
            logger.info('数据库连接关闭!')
        except Exception as e:
            logger.error(f'关闭数据库连接失败，报错信息：{e}')
